import datetime
import logging
import os
import sys
import trio
from contextlib import AsyncExitStack
from functools import partial
from typing import Optional, List

import baostock as bs
import click
import MyTT
import numpy as np
import pandas as pd
from email.mime.text import MIMEText
from email.header import Header
import smtplib
from sqlalchemy import create_engine, text, MetaData, Table
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from models import Base

# 初始化日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 数据库连接
DATABASE_URL = "mysql+asyncmy://taomu:taomu-rcmu#2025@192.168.1.220:3306/taomu_stock"
engine = create_async_engine(DATABASE_URL, echo=False)

# 同步引擎用于删除表等操作
sync_engine = create_engine("mysql+pymysql://taomu:taomu-rcmu#2025@192.168.1.220:3306/taomu_stock")
Base.metadata.create_all(sync_engine)

# 脚本开始时间
SCRIPT_START_TIME = datetime.datetime.now()


def get_stock_history_data(dir,code,exists='append'):
    sdate = (datetime.datetime.today()-datetime.timedelta(days=600)).strftime("%Y-%m-%d")
    print(sdate,code)
    rs = bs.query_history_k_data_plus(code,"date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST",start_date=sdate,end_date=None,frequency="d")
    data_list=[]
    while (rs.error_code == '0') & rs.next():
        row_data = rs.get_row_data()
        if row_data[-1] == "0":
            data_list.append(row_data)
    else:
        result = pd.DataFrame(data_list,columns=rs.fields)
        result['date'] = pd.to_datetime(result['date']);
        result['close'] = result['close'].astype(float);
        result['high'] = result['high'].astype(float);
        result['low'] = result['low'].astype(float);
        result['open'] = result['open'].astype(float);
        result['code'] = result['code'].astype(str)
        # 替换空字符串为 NaN
        result['volume'] = result['volume'].replace('', np.nan)
        # 转换为浮点数类型
        result['volume'] = pd.to_numeric(result['volume'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['volume'] = result['volume'].fillna(0)
        result['volume'] = result['volume'].astype(float);
        result['preclose'] = result['preclose'].astype(float);
        # 替换空字符串为 NaN
        result['amount'] = result['amount'].replace('', np.nan)
        # 转换为浮点数类型
        result['amount'] = pd.to_numeric(result['amount'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['amount'] = result['amount'].fillna(0)
        result['amount'] = result['amount'].astype(float); # 替换空字符串为 NaN
        result['turn'] = result['turn'].replace('', np.nan)
        # 转换为浮点数类型
        result['turn'] = pd.to_numeric(result['turn'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['turn'] = result['turn'].fillna(0)
        result['turn'] = result['turn'].astype(float);
         # 替换空字符串为 NaN
        result['pctChg'] = result['pctChg'].replace('', np.nan)
        # 转换为浮点数类型
        result['pctChg'] = pd.to_numeric(result['pctChg'], errors='coerce')
        # 填充缺失值为 0（可选）
        result['pctChg'] = result['pctChg'].fillna(0)
        result['pctChg'] = result['pctChg'].astype(float);
        #  使用MyTT计算指标 5日均线、10日均线、20日均线、30日均线、60日均线 添加到result中 
        closes = result['close']
        result['obv'] = MyTT.OBV(closes,result['volume'])
        result['M5'] = MyTT.MA(closes,5)
        result['M10'] = MyTT.MA(closes,10)
        result['M20'] = MyTT.MA(closes,20)
        result['M30'] = MyTT.MA(closes,30)
        result['M50'] = MyTT.MA(closes,50)
        result['M60'] = MyTT.MA(closes,60)
        result['EMA5'] = MyTT.EMA(closes,5)
        result['EMA10'] = MyTT.EMA(closes,10)
        result['EMA20'] = MyTT.EMA(closes,20)
        result['BBI'] = MyTT.BBI(closes)
        result['boll_upper'], result['boll_middle'], result['boll_lower'] = MyTT.BOLL(closes)
        result['macd_dif'], result['macd_dea'], result['macd'] = MyTT.MACD(closes)
        result['rsi6'] = MyTT.RSI(closes,6)
        result['rsi12'] = MyTT.RSI(closes,12)
        result['k'],result['d'],result['j'] = MyTT.KDJ(result['high'],result['low'],closes)

        # 去掉开头60天数据
        result = result[result.index > 60]
        result.to_sql("taomu_stock_history", sync_engine, if_exists="append", index=False)


def get_stock_code_list():
    d = 1
    df = bs.query_all_stock(day=datetime.date.today() - datetime.timedelta(days=d)).get_data()
    while len(df) == 0:
        d += 1
        df = bs.query_all_stock(day=datetime.date.today() - datetime.timedelta(days=d)).get_data()
    print(len(df))
    if len(df) > 0:
        codes_df = df[
            ((df['code'] < 'sz.300000') & (df['code'] > 'sh.689000') |
             (df['code'] < 'sh.688000') & (df['code'] > 'sh.009999')) &
            (~df['code_name'].str.startswith('ST')) &
            (~df['code_name'].str.startswith('*ST'))
        ]
        return codes_df

@click.command()
@click.option("--dir", prompt="csv存放的目录", help="csv存放的目录")
def main(dir):
    os.makedirs(dir, exist_ok=True)
    bs.login()

    # 获取股票列表
    cpd = get_stock_code_list()
    clist = cpd['code'].to_list()
    stock_count = len(clist)

    # 清空表
    with sync_engine.connect() as connection:
        try:
            connection.execute(text("delete from taomu_stock_list;"))
        except:
            pass
        try:
            connection.execute(text("delete from taomu_stock_history;"))
        except:
            pass
        connection.commit()

    # 插入股票列表
    cpd.to_sql("taomu_stock_list", sync_engine, if_exists='append', index=False)

    get_stock_history_data(dir,clist.pop(),'replace')
    for code in clist:
        get_stock_history_data(dir,code)
    #get_stock_history_data("sh.600838")

    bs.logout()

    # 发送邮件报告
    end_time = datetime.datetime.now()
    send_email_report(SCRIPT_START_TIME, end_time, stock_count, count_total_rows())

    logger.info("全部完成")


def count_total_rows():
    with sync_engine.connect() as conn:
        return conn.execute(text("SELECT COUNT(*) FROM taomu_stock_history")).fetchone()[0]


def send_email_report(start_time, end_time, stock_count, history_count):
    smtp_server = "smtp.126.com"
    smtp_port = 465
    sender = "minix3@126.com"
    password = "LAZpFpGT9zHb6C3f"
    receiver = "minix3@126.com"

    duration = (end_time - start_time).total_seconds()

    content = f"""
    股票数据抓取完成，请查收本次执行结果：

    开始时间: {start_time}
    结束时间: {end_time}
    总耗时: {int(duration)} 秒
    抓取股票总数: {stock_count} 只
    历史数据总条数: {history_count} 条

    —— 自动化脚本 by Qwen
    """

    message = MIMEText(content, 'plain', 'utf-8')
    message['From'] = Header("股票抓取系统", 'utf-8')
    message['To'] = Header("管理员", 'utf-8')
    message['Subject'] = Header("[自动报告] 股票历史数据抓取完成", 'utf-8')

    try:
        server = smtplib.SMTP_SSL(smtp_server, smtp_port)
        server.login(sender, password)
        server.sendmail(sender, [receiver], message.as_string())
        print("邮件发送成功")
    except Exception as e:
        print("邮件发送失败:", str(e))
    finally:
        server.quit()


if __name__ == '__main__':
    main()
